import numpy as np
from scipy.signal import savgol_filter


class RealTimeSavitzkyGolayFilter:
    def __init__(self, window_length=51, polyorder=3):
        """
        初始化Savitzky-Golay滤波器
        :param window_length: 窗口长度，必须为正奇数
        :param polyorder: 多项式阶数，必须小于窗口长度
        """
        self.window_length = window_length
        self.polyorder = polyorder
        self.buffer = []

    def add_sample(self, sample):
        """
        添加新的样本到缓冲区
        :param sample: 新的样本值
        """
        self.buffer.append(sample)
        if len(self.buffer) > self.window_length:
            self.buffer.pop(0)

    def get_filtered_value(self):
        """
        获取滤波后的最新值
        :return: 滤波后的值
        """
        if len(self.buffer) < self.window_length:
            return self.buffer[-1]  # 如果缓冲区未满，返回最后一个值
        filtered_values = savgol_filter(self.buffer, self.window_length, self.polyorder)
        return filtered_values[-1]


# 示例使用
if __name__ == "__main__":
    # 假设我们有一个实时的sEMG数据流
    semg_data_stream = [np.random.rand() * 1000 for _ in range(1000)]  # 模拟sEMG数据

    # 创建Savitzky-Golay滤波器实例
    sg_filter = RealTimeSavitzkyGolayFilter(window_length=51, polyorder=3)

    # 处理实时数据流
    filtered_data = []
    for sample in semg_data_stream:
        sg_filter.add_sample(sample)
        filtered_value = sg_filter.get_filtered_value()
        filtered_data.append(filtered_value)

    # 打印原始数据和滤波后的数据
    import matplotlib.pyplot as plt

    plt.figure(figsize=(12, 6))
    plt.plot(semg_data_stream, label='原始sEMG数据')
    plt.plot(filtered_data, label='滤波后的sEMG数据', linewidth=2)
    plt.legend()
    plt.show()
